import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.colors import LogNorm

fname = 'fig2b.npz'

# load npz file
data = np.load(fname, allow_pickle=True)
print(data.files)
n = data['n']
D = data['D']
nu = data['nu']
nH = data['nH']
Vtg = data['Vtg']
Vbg = data['Vbg']
n0 = data['n0'] # density at charge neutrality point
D0 = data['D0'] 
ns = data['ns'] # density at 4 electron per moire unit cell
Rxy1 = data['Rxy1']
Rxy2 = data['Rxy2']
B1 = data['B1']
B2 = data['B2']
device_name = data['device_name']
back_gate_thickness_nm = data['back_gate_thickness_nm']
top_gate_thickness_nm = data['top_gate_thickness_nm']

# plot
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
print('n.shape = ', n.shape)
print('D.shape = ', D.shape)
print('nH.shape = ', nH.shape)
im = ax.pcolormesh(n, D, nH, cmap='seismic', vmin=-10, vmax=10)
ax.set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax.set_ylabel('$D/\epsilon_0$ (V/nm)')
ax.set_title('Fig. 2b')
cbar = fig.colorbar(im, ax=ax)
cbar.set_label('$n_\mathrm{H}$ ($10^{12}$ cm$^{-2}$)')
plt.show()
